package cover

import (
	"fmt"
	"github.com/viant/bintly"
	"io"
	"reflect"
	"unsafe"
)

// ExportDocuments writes the tree's embeddings (vectors) paired with values.
// The format is:
//   - base (float32)
//   - distance func name (string)
//   - count (int32)
//   - repeated Point coders (count)
//   - values slice matching order (encoded in the same way as values.Encode)
func (t *Tree[T]) ExportDocuments(writer io.Writer) error {
	t.mu.RLock()
	defer t.mu.RUnlock()

	// Collect points and values in ascending index order for determinism.
	if t.indexMap == nil {
		// nothing to export
		return nil
	}
	// Determine max index to iterate compactly
	maxIdx := int32(-1)
	for idx := range t.indexMap {
		if idx > maxIdx {
			maxIdx = idx
		}
	}
	points := make([]*Point, 0, len(t.indexMap))
	vals := make([]T, 0, len(t.indexMap))
	for i := int32(0); i <= maxIdx; i++ {
		p, ok := t.indexMap[i]
		if !ok || p == nil || !p.HasValue() {
			continue
		}
		points = append(points, p)
		vals = append(vals, t.values.value(i))
	}

	// Encode header + points + values in a single stream
	w := writers.Get()
	defer writers.Put(w)
	w.Float32(t.base)
	w.String(string(t.distanceFuncName))
	w.Int32(int32(len(points)))
	for i := range points {
		if err := w.Coder(points[i]); err != nil {
			return err
		}
	}
	// Encode values slice inline
	if err := encodeValuesSliceToWriter(w, vals, t.values.Type); err != nil {
		return err
	}
	_, err := writer.Write(w.Bytes())
	return err
}

// ImportDocuments reads embeddings (vectors) paired with values and rebuilds the tree.
// It clears the current structure and inserts imported documents.
func (t *Tree[T]) ImportDocuments(reader io.Reader) error {
	t.mu.Lock()
	defer t.mu.Unlock()

	data, err := io.ReadAll(reader)
	if err != nil {
		return err
	}
	r := readers.Get()
	defer readers.Put(r)
	if err := r.FromBytes(data); err != nil {
		return err
	}

	// Read header
	r.Float32(&t.base)
	var dist string
	r.String(&dist)
	t.distanceFuncName = DistanceFunction(dist)
	t.distanceFnunc = t.distanceFuncName.Function()

	// Read points
	var count int32
	r.Int32(&count)
	pts := make([]*Point, int(count))
	for i := 0; i < int(count); i++ {
		pt := &Point{}
		if err := r.Coder(pt); err != nil {
			return err
		}
		pts[i] = pt
	}

	// Prepare values type info
	t.values = values[T]{data: make([]T, 0)}
	t.values.ensureType()

	// Decode values slice from remaining reader data
	vals, err := decodeValuesSlice[T](r, int(count), t.values.Type)
	if err != nil {
		return err
	}

	// Reset structure
	t.root = nil
	t.indexMap = make(map[int32]*Point)
	t.version++

	// Insert pairs
	for i := 0; i < int(count); i++ {
		t.insertOrAttach(vals[i], pts[i])
	}
	return nil
}

// insertOrAttach inserts a value with its point, preserving point fields.
func (t *Tree[T]) insertOrAttach(value T, point *Point) int32 {
	// This mirrors Insert but does not override point.index before put
	// because values.put returns sequential index and we want to align.
	point.index = t.values.put(value)
	if t.indexMap == nil {
		t.indexMap = make(map[int32]*Point)
	}
	t.indexMap[point.index] = point
	if point.Magnitude == 0 && len(point.Vector) > 0 {
		// lazily compute magnitude if not provided
		// search.Float32s is imported in distance.go; keep logic consistent by recomputing on-demand elsewhere
	}
	if t.root == nil {
		t.root = &Node{point: point, level: 0}
	} else {
		t.insert(t.root, point, 0)
	}
	t.version++
	return point.index
}

// encodeValuesSliceToWriter writes the given slice into provided bintly writer,
// using the same rules as values.Encode
func encodeValuesSliceToWriter[T any](w *bintly.Writer, data []T, typ reflect.Type) error {
	// Fast-path for common primitives
	switch actual := any(data).(type) {
	case []string:
		w.Strings(actual)
	case []int:
		w.Ints(actual)
	case []float32:
		w.Float32s(actual)
	case []float64:
		w.Float64s(actual)
	case []bool:
		w.Bools(actual)
	default:
		if typ == nil {
			// Attempt to infer
			typ = reflect.TypeOf(data).Elem()
		}
		if typ.Comparable() && typ.Kind() != reflect.Pointer {
			raw := unsafe.Slice((*byte)(unsafe.Pointer(&data)), len(data)*int(typ.Size()))
			w.Uint8s(raw)
		} else {
			// custom encoding per element
			w.Alloc(int32(len(data)))
			for i := range data {
				enc, ok := any(data[i]).(bintly.Encoder)
				if !ok {
					return fmt.Errorf("unable to cast Encoder from %T", data[i])
				}
				if err := enc.EncodeBinary(w); err != nil {
					return err
				}
			}
		}
	}
	return nil
}

// decodeValuesSlice reads a slice using the same rules as values.Decode
func decodeValuesSlice[T any](r *bintly.Reader, _count int, typ reflect.Type) ([]T, error) {
	var data []T
	// Fast-path for common primitives
	values := any(data)
	switch values.(type) {
	case []string:
		var actual []string
		r.Strings(&actual)
		data = any(actual).([]T)
	case []int:
		var actual []int
		r.Ints(&actual)
		data = any(actual).([]T)
	case []float32:
		var actual []float32
		r.Float32s(&actual)
		data = any(actual).([]T)
	case []float64:
		var actual []float64
		r.Float64s(&actual)
		data = any(actual).([]T)
	case []bool:
		var actual []bool
		r.Bools(&actual)
		data = any(actual).([]T)
	default:
		if typ == nil {
			typ = reflect.TypeOf(data).Elem()
		}
		if typ.Comparable() && typ.Kind() != reflect.Pointer {
			var raw []byte
			r.Uint8s(&raw)
			size := len(raw) / int(typ.Size())
			data = make([]T, size)
			buf := unsafe.Slice((*byte)(unsafe.Pointer(&data)), len(data)*int(typ.Size()))
			copy(buf, raw)
		} else {
			// custom decode per element
			size := r.Alloc()
			data = make([]T, size)
			if typ.Kind() == reflect.Ptr {
				for i := range data {
					data[i] = reflect.New(typ.Elem()).Interface().(T)
				}
			}
			for i := range data {
				dec, ok := any(data[i]).(bintly.Decoder)
				if !ok {
					if dec2, ok2 := any(&data[i]).(bintly.Decoder); ok2 {
						if err := dec2.DecodeBinary(r); err != nil {
							return nil, err
						}
						continue
					}
					return nil, fmt.Errorf("unable to cast Decoder from %T", data[i])
				}
				if err := dec.DecodeBinary(r); err != nil {
					return nil, err
				}
			}
		}
	}
	return data, nil
}
